package com.mr.task;

import com.mr.core.*;
import com.mr.job.Job;
import com.mr.split.InputSplit;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;
import java.io.ObjectOutputStream;
import java.io.FileOutputStream;
import java.io.IOException;

public class MapTask<K1, V1, K2, V2> implements Callable<Map<Integer, String>> {
    private final InputSplit split;
    private final int taskId;
    private final Job<K1, V1, K2, V2, ?, ?> job;

    public MapTask(InputSplit split, Job<K1, V1, K2, V2, ?, ?> job, int taskId) {
        this.split = split;
        this.job = job;
        this.taskId = taskId;
    }

    @Override
    public Map<Integer, String> call() throws Exception {
        try {
            // 调试日志
            System.out.println("Starting map task " + taskId + " for file: " + split.getFile());
            
            Mapper<K1, V1, K2, V2> mapper = job.getMapperClass().getDeclaredConstructor().newInstance();
            Context<K2, V2> context = new Context<>();

            // 使用split的readLines方法获取输入数据
            List<String> lines = split.readLines();
            for (String line : lines) {
                try {
                    @SuppressWarnings("unchecked")
                    K1 key = (K1) line;
                    @SuppressWarnings("unchecked")
                    V1 value = (V1) line;
                    mapper.map(key, value, context);
                } catch (Exception e) {
                    context.reportError(e);
                }
            }

            // Combiner实例化
            if (job.getCombinerClass() != null) {
                Combiner<K2, V2> combiner = job.getCombinerClass().getDeclaredConstructor().newInstance();
                Map<Integer, Map<K2, List<V2>>> partitioned = context.getPartitioned(job.getNumReduceTasks());
                for (Map.Entry<Integer, Map<K2, List<V2>>> entry : partitioned.entrySet()) {
                    Context<K2, V2> combinerContext = new Context<>();
                    entry.getValue().forEach((k, v) -> 
                        combiner.combine(k, v, combinerContext));
                    entry.setValue(combinerContext.getBuffer());
                }
            }

            // 检查错误状态
            if (context.hasError()) {
                throw new RuntimeException("Map task failed", context.getError());
            }

            // 写入中间文件
            Map<Integer, String> partitionFiles = new HashMap<>();
            Map<Integer, Map<K2, List<V2>>> partitioned = context.getPartitioned(job.getNumReduceTasks());
            for (Map.Entry<Integer, Map<K2, List<V2>>> entry : partitioned.entrySet()) {
                String filename = String.format("%s/map-%d-part-%d.tmp", 
                    job.getIntermediateDir(), taskId, entry.getKey());
                writeToFile(entry.getValue(), filename);
                partitionFiles.put(entry.getKey(), filename);
            }
            return partitionFiles;
        } catch (Exception e) {
            throw new RuntimeException("Map task execution failed", e);
        }
    }

    private void writeToFile(Map<K2, List<V2>> data, String filename) {
        try (ObjectOutputStream oos = new ObjectOutputStream(new FileOutputStream(filename))) {
            oos.writeObject(data);
        } catch (IOException e) {
            throw new RuntimeException("Failed to write map output", e);
        }
    }
}